import numpy as np
from ch06.rnnlm import Rnnlm
from common.functions import softmax


class RnnlmGen(Rnnlm):
    def generate(self,start_id,skip_ids=None,sample_size=100):
        word_ids = [start_id] 

        x = start_id
        while len(word_ids) < sample_size:
            x = np.array(x).reshape(1,1)
            score = self.predict(x)
            p = softmax(score.flatten())

            sampled = np.random.choice(len(p),size=1,p=p)
            if (skip_ids is None) or (sampled not in skip_ids):
                x = sampled
                word_ids.append(int(x))

        return word_ids
    
    def get_state(self):
        return self.lstm_layer.h, self.lstm_layer.c

    def set_state(self, state):
        self.lstm_layer.set_state(*state)